#if 0
Copyright 2011 Google. All rights reserved.

Redistribution and use in source and binary forms, with or without modification, are
permitted provided that the following conditions are met:

   1. Redistributions of source code must retain the above copyright notice, this list of
      conditions and the following disclaimer.

   2. Redistributions in binary form must reproduce the above copyright notice, this list
      of conditions and the following disclaimer in the documentation and/or other materials
      provided with the distribution.

THIS SOFTWARE IS PROVIDED BY <COPYRIGHT HOLDER> ''AS IS'' AND ANY EXPRESS OR IMPLIED
WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL <COPYRIGHT HOLDER> OR
CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

The views and conclusions contained in the software and documentation are those of the
authors and should not be interpreted as representing official policies, either expressed
or implied, of Google.
#endif

#include "win32_allocator_patcher.h"

#include <assert.h>
#include <iostream>
#include <malloc.h>
#include <new>
#include <stdio.h>

#define WIN32_LEAN_AND_MEAN 
#include "windows.h"
#include "ms_detours/detours.h"

#include "win32_allocator_interface.h"

#ifndef DCHECK
# define DCHECK(X) assert((X));
#endif

namespace alloc_override {

///////////////////////////////////////////////////////////////////////////////
static Win32AllocatorInterface* s_allocator_override_ = NULL;
///////////////////////////////////////////////////////////////////////////////

// constructs the RawWin32SystemAllocator on startup so that it is immediately
// available.
#pragma optimize("", off)
// turn off optimizations to prevent linker from stripping unused symbol. This
// object is here soley for the purpose of constructing the allocator.
namespace alloc_override_detail {
Win32AllocatorInterface& on_startup = RawWin32SystemAllocator();
}
#pragma optimize("", on)


typedef void* (* FunMalloc)(size_t sz);
typedef void* (* FunCalloc)(size_t num, size_t sz);
typedef void (* FunFree)(void* ptr);
typedef void* (* FunRealloc)(void* ptr, size_t sz);
typedef void* (* FunExpand)(void* memblock, size_t sz);
typedef size_t (* FunMsize)(void* memblock);
typedef char* (*FunStrdup) (const char* str);

typedef void* (* FunNew)(std::size_t size);
typedef void* (* FunNewNoThrow)(std::size_t size,
                                const std::nothrow_t& nothrow_constant);
typedef void* (* FunNewArray)(std::size_t size);
typedef void* (* FunNewArrayNoThrow)(std::size_t size,
                                     const std::nothrow_t& nothrow_constant);
typedef void (* FunDelete)(void* ptr);
typedef void (* FunDeleteArray)(void* ptr);

// s_real_XXX will always point to the real versions of the allocators. This
// is by design of the MS Detours program.
static FunMalloc s_real_malloc = ::malloc;
static FunCalloc s_real_calloc = ::calloc;
static FunFree s_real_free = ::free;
static FunRealloc s_real_realloc = ::realloc;
static FunExpand s_real_expand = ::_expand;
static FunMsize s_real_msize = ::_msize;
static FunNew s_real_new = ::operator new;
static FunNewNoThrow s_real_new_no_throw = ::operator new;
static FunNewArray s_real_new_array = ::operator new[];
static FunNewArrayNoThrow s_real_new_array_no_throw = ::operator new[];
static FunDelete s_real_delete = ::operator delete;
static FunDeleteArray s_real_delete_array = ::operator delete[];
static FunStrdup s_real_strdup = _strdup;

class RawWin32Allocator : public Win32AllocatorInterface {
 public:
  virtual void* Malloc(size_t sz) {
    return s_real_malloc(sz);
  }
  virtual void* Calloc(size_t num, size_t sz) {
    return s_real_calloc(num, sz);
  }
  virtual void* Realloc(void* ptr, size_t sz) {
    return s_real_realloc(ptr, sz);
  }
  virtual void* Expand(void* memblock, size_t sz) {
    return s_real_expand(memblock, sz);
  }
  virtual char* Strdup(const char* str) {
    return s_real_strdup(str);
  }
  virtual size_t Msize(void* memblock) {
    return s_real_msize(memblock);
  }
  virtual void Free(void* ptr) {
    s_real_free(ptr);
  }

  virtual void* New(std::size_t size) {
    return s_real_new(size);
  }
  virtual void* NewNoThrow(std::size_t size) {
    return s_real_new_no_throw(size, std::nothrow_t());
  }
  virtual void* NewArray(std::size_t size) {
    return s_real_new_array(size);
  }
  virtual void* NewArrayNoThrow(std::size_t size) {
    return s_real_new_array_no_throw(size, std::nothrow_t());
  }
  virtual void Delete(void* ptr) {
    s_real_delete(ptr);
  }
  virtual void DeleteArray(void* ptr) {
    s_real_delete_array(ptr);
  }
};

static void* MyMalloc(size_t sz) {
  return s_allocator_override_->Malloc(sz);
}

static void* MyCalloc(size_t num, size_t sz) {
  return s_allocator_override_->Calloc(num, sz);
}

static void MyFree(void* ptr) {
  s_allocator_override_->Free(ptr);
}

static void* MyRealloc(void* ptr, size_t sz) {
  return s_allocator_override_->Realloc(ptr, sz);
}

static void* MyExpand(void* ptr, size_t sz) {
  return s_allocator_override_->Expand(ptr, sz);
}

static size_t MyMsize(void* memblock) {
  return s_allocator_override_->Msize(memblock);
}

static void* MyNew(std::size_t sz) {
  return s_allocator_override_->New(sz);
}

static void* MyNewNoThrow(std::size_t sz,
                          const std::nothrow_t& no_exception) {
  return s_allocator_override_->NewNoThrow(sz);
}

static void* MyNewArray(std::size_t sz) {
  return s_allocator_override_->NewArray(sz);
}

static void* MyNewArrayNoThrow(std::size_t sz,
                               const std::nothrow_t& no_exception) {
  return s_allocator_override_->NewArrayNoThrow(sz);
}

static void MyDelete(void* ptr) {
  s_allocator_override_->Delete(ptr);
}

static void MyDeleteArray(void* ptr) {
  s_allocator_override_->DeleteArray(ptr);
}

static char* MyFunStrdup(const char* str) {
  return s_allocator_override_->Strdup(str);
}

///////////////////////////////////////////////////////////////////////////////
void OverrideAllocatorOnce(Win32AllocatorInterface* allocator) {
  static bool init_once = false;
  if (init_once) {
    DCHECK(false);
    return;
  }
  init_once = true;
  s_allocator_override_ = allocator;
  // Stops all threads except main.
  DetourTransactionBegin();
  DetourUpdateThread(GetCurrentThread());
  // Hot patch all the system allocator functions in the current running
  // dll.
  DetourAttach(&(PVOID&)s_real_malloc, MyMalloc);
  DetourAttach(&(PVOID&)s_real_calloc, MyCalloc);
  DetourAttach(&(PVOID&)s_real_free, MyFree);
  DetourAttach(&(PVOID&)s_real_realloc, MyRealloc);
  DetourAttach(&(PVOID&)s_real_expand, MyExpand);
  DetourAttach(&(PVOID&)s_real_msize, MyMsize);
  DetourAttach(&(PVOID&)s_real_new, MyNew);
  DetourAttach(&(PVOID&)s_real_new_no_throw, MyNewNoThrow);
  DetourAttach(&(PVOID&)s_real_new_array, MyNewArray);
  DetourAttach(&(PVOID&)s_real_new_array_no_throw, MyNewArrayNoThrow);
  DetourAttach(&(PVOID&)s_real_delete, MyDelete);
  DetourAttach(&(PVOID&)s_real_delete_array, MyDeleteArray);
  DetourAttach(&(PVOID&)s_real_strdup, MyFunStrdup);
  // Purges the instruction cache so that all threads will see the updated
  // versions of the allocators and then resumes all threads.
  DetourTransactionCommit();
}

///////////////////////////////////////////////////////////////////////////////
void* AlignBytes(int multiple, void* input) {
  intptr_t ptr_bits = reinterpret_cast<intptr_t>(input);
  const intptr_t kLowOrderBits = multiple - 1;
  ptr_bits += kLowOrderBits;
  ptr_bits &= ~kLowOrderBits;
  return reinterpret_cast<void*>(ptr_bits);
}

///////////////////////////////////////////////////////////////////////////////
Win32AllocatorInterface& RawWin32SystemAllocator() {
  // This funky business here does two things:
  // 1) The returned allocator uses static memory to allocate.
  // 2) The returned allocator will never run it's destructor.
  static const int kAlign = 8;
  static char* memory[kAlign + sizeof(RawWin32Allocator)];
  static bool initialized = false;
  static RawWin32Allocator* ptr = NULL;
  if (!initialized) {
    initialized = true;
    // Use 8-byte alignment.
    ptr = static_cast<RawWin32Allocator*>(AlignBytes(8, memory));
    // Construct the allocator at the memory position. This version of new
    // is not overriden at all.
    new (ptr) RawWin32Allocator;
  }
  DCHECK(ptr != NULL);
  return *ptr;
}
}  // namespace alloc_override